import cv2
import math
import numpy as np


def generate_video(path):
    """

    :param path:
    :return:
    """
    fps = 30
    # 图片尺寸
    img_size = (640, 480)
    fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
    # face_video_writer = cv2.VideoWriter(face_path, fourcc, fps, img_size)
    cap = cv2.VideoCapture(path)
    number = cap.get(7)
    print("width:{}, height:{}, frame number{}".format(
        cap.get(cv2.CAP_PROP_FRAME_WIDTH),
        cap.get(cv2.CAP_PROP_FRAME_HEIGHT),
        number)
    )
    if number >= 180:
        i = 0
        frames = []
        while True:
            i += 1
            ret, img = cap.read()
            if not ret:
                break
            # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            frames.append(img)
        length = math.floor((number / 180) * 2 - 1)
        split = str(path).split("\\")[-1]
        split = split.split("-")
        root_path = "D:\\dataset\\YawDDdata_original\\Mirror\\all_cut\\"
        for i in range(0, length):
            cut_path = root_path + split[0] + "-" + split[1] + "_" + str(i) + "-" + split[2]
            print(cut_path)
            face_video_writer = cv2.VideoWriter(cut_path, fourcc, fps, img_size)
            for im_f in frames[i * 60:i * 60 + 180]:
                face_video_writer.write(im_f)
            face_video_writer.release()
    cap.release()
    print('finish')
    pass


def generate_video_enhancement(path):
    """

        :param path:
        :return:
        """
    fps = 30
    # img_size = (640, 480)
    img_size = (512, 384)
    fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
    cap = cv2.VideoCapture(path)
    number = cap.get(7)
    print("path:{} width:{}, height:{}, frame number{}".format(
        path,
        cap.get(cv2.CAP_PROP_FRAME_WIDTH),
        cap.get(cv2.CAP_PROP_FRAME_HEIGHT),
        number)
    )
    if number >= 180:
        i = 0
        frames = []
        while True:
            i += 1
            ret, img = cap.read()
            if not ret:
                break
            # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            frames.append(img)
        frames = np.array(frames)
        videos = [
            frames[:, :384, :512],
            frames[:, 96:, :512],
            frames[:, 48:-48, 64:-64],
            frames[:, :384, 128:],
            frames[::, 96:, 128:],
        ]
        video_name = ["top_left", "top_right", "center", "bottom_left", "bottom_right"]
        split = str(path).split("\\")[-1]
        split = split.split("-")
        root_path = "D:\\dataset\\YawDDdata_original\\Mirror\\all_cut_enhancement\\"
        for i in range(len(videos)):
            cut_path = root_path + split[0] + "-" + split[1] + "_" + video_name[i] + "-" + split[2]
            print(cut_path)
            face_video_writer = cv2.VideoWriter(cut_path, fourcc, fps, img_size)
            for im_f in videos[i]:
                face_video_writer.write(im_f)
            face_video_writer.release()
    print('finish')


def number_of_frame(path):
    cap = cv2.VideoCapture(path)
    number = cap.get(7)
    cap.release()
    return number
